import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.feature_selection import mutual_info_regression


plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def Mutual_feature_label(x, y, logger):
    logger.info('--------输出特征和标签之间的相关性---------')
    # 使用互信息（Mutual Information, MI）来检索特征和标签之间的相关性
    average_mi = pd.Series(0, index=x.columns, dtype='float64')
    for label in y:
        mi = mutual_info_regression(x, y[label])
        mi_scores = pd.Series(mi, index=x.columns).sort_values(ascending=False)
        logger.info(f"{label} Mutual Information:\n", mi_scores)
        average_mi += mi_scores
    average_mi /= 3
    average_mi_scores = pd.Series(average_mi, index=x.columns).sort_values(ascending=False)
    logger.info("平均 Mutual Information:\n", average_mi_scores)


def feature_correlation(x, logger):
    # 计算相关系数矩阵，绘制热力图，来查看特征之间的相关性
    logger.info('--------绘制特征之间相关性的热力图---------')
    corr_matrix = x.corr()
    plt.figure(figsize=(10, 10))
    sns.heatmap(corr_matrix, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
    plt.title("Feature Correlation Heatmap")
    plt.xticks(rotation=60)
    plt.tight_layout()
    # os.makedirs('feature_img', exist_ok=True)
    # img_feature = "./feature_img/feature_corr_heatmap_" + sequence + '.png'
    # plt.savefig(img_feature)
    plt.show()


def label_correlation(y, logger):
    # 计算相关系数矩阵，绘制热力图，来查看相关性
    logger.info('--------绘制标签之间相关性的热力图---------')
    label_corr = y[['资产总额增长率', '主营业务收入增长率', '利润总额增长率']].corr()
    plt.figure(figsize=(10, 10))
    sns.heatmap(label_corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1)
    plt.title("Label Correlation Heatmap")
    # img_label = './feature_img/label_corr_heatmap_' + sequence + '.png'
    # plt.savefig(img_label)
    plt.show()


def feature_engineer_origin(data, logger):
    x = data.iloc[:, :-3]
    y = data.iloc[:, -3:]
    feature_names = ['净资产（万元）', '成本费用总额（万元）', '销售收入（万元）', '研发费用总额（万元）', '纳税总额（万元）',
                     '企业职工总数', '企业科技人员总数', '总知识产权数']
    x = x[feature_names]

    Mutual_feature_label(x, y, logger)
    feature_correlation(x, logger)
    label_correlation(y, logger)

    return x, y


def feature_engineer(data, logger):
    logger.info('--------开始特征工程---------')

    x = data.iloc[:, :-3]
    y = data.iloc[:, -3:]

    x['人均销售收入'] = x['销售收入（万元）'] / x['企业职工总数']
    x['科技人员占比'] = x['企业科技人员总数'] / x['企业职工总数']

    # 2. 对数转换可能有助于非线性关系和偏态分布
    for col in ['销售收入（万元）']:
        x[f'log_{col}'] = np.log1p(x[col])

    Mutual_feature_label(x, y, logger)
    feature_correlation(x, logger)

    feature_names = ['研发费用占销售收入总额比例', '研发费用占成本费用支出总额比例', '净资产（万元）',

                     '人均销售收入', 'log_销售收入（万元）',
                     ]
    x = x[feature_names]

    feature_correlation(x, logger)

    return x, y
